package cn.itcast.tags.models.statistics

import cn.itcast.tags.models.{AbstractModel, ModelType}
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.types.StringType

class AgeRangeModel extends AbstractModel("AgeRangeModel", ModelType.STATISTICS) {
  override def doTag(businessDF: DataFrame, tagDF: DataFrame): DataFrame = {
    import businessDF.sparkSession.implicits._
    import org.apache.spark.sql.functions._

    val rule_to_tuple: UserDefinedFunction = udf(
      (rule: String) => {
        val Array(start, end) = rule.trim.split("-").map(_.toInt)
        (start, end)
      }
    )
    val attrTagDF: DataFrame = tagDF.filter($"level" === 5)
      .select(
        $"id".as("tagId"),
        rule_to_tuple($"rule").as("rules")
      )
      .select(
        $"tagId",
        $"rules._1".as("start"),
        $"rules._2".as("end")
      )

    val birthdayDF: DataFrame = businessDF
      .select(
        $"id".as("uid"),
        regexp_replace($"birthday", "-", "").cast(StringType).as("bornDate")
      )

    val modelDF: DataFrame = birthdayDF.join(attrTagDF)
      .where($"bornDate".between($"start", $"end"))
      .select(
        $"uid",
        $"tagId".cast(StringType)
      )


    modelDF
  }
}

object AgeRangeModel {
  def main(args: Array[String]): Unit = {
    val tagModel = new AgeRangeModel
    tagModel.executeModel(332L, false)
  }
}


